﻿using System;
using System.Collections.Generic;

namespace Neo4jConnect
{
    public static class NodeExtensions
    {
        public static IEnumerable<Relationship> GetRelationships(this Node node, RelationshipDirection direction)
        {
            foreach (var relationship in node.GetRelationships())
            {
                if (direction == RelationshipDirection.In && relationship.StartUri != node.SelfUri)
                {
                    yield return relationship;
                }
                else if (direction == RelationshipDirection.Out && relationship.StartUri == node.SelfUri)
                {
                    yield return relationship;
                }
                else if (direction == RelationshipDirection.All)
                {
                    yield return relationship;
                }
            }
        }

        public static IEnumerable<Relationship> GetRelationships(this Node node, Enum relationshipType)
        {
            string reltype = relationshipType.ToString();
            foreach (var relationship in node.GetRelationships())
            {
                if (relationship.Type == reltype)
                {
                    yield return relationship;
                }
            }
        }

        public static IEnumerable<Relationship> GetRelationships(this Node node, RelationshipDirection direction, Enum relationshipType)
        {
            string reltype = relationshipType.ToString();
            foreach (var relationship in node.GetRelationships())
            {
                if (relationship.Type == reltype)
                {
                    if (direction == RelationshipDirection.In && relationship.StartUri == node.SelfUri)
                    {
                        yield return relationship;
                    }
                    else if (direction == RelationshipDirection.Out && relationship.StartUri != node.SelfUri)
                    {
                        yield return relationship;
                    }
                    else if (direction == RelationshipDirection.All)
                    {
                        yield return relationship;
                    }
                }
            }
        } 
    }
}
